BioJulia系列-Automa

fasta_regex = let
    header = re"[a-z]+"
    seqline = re"[ACGTU]+"
    record = '>' * header * '\n' * rep1(seqline) * `\n`
    rep(record)
end;

julia> eval(generate_buffer_validator(:validate_fasta, fasta_regex));

julia> validate_fasta
validate_fasta (generic function with 1 method)

julia> validate_fasta(">hello\nTAGAGA\nTAGAG") # missing trailing newline
0

julia> validate_fasta(">helloXXX") # Error at byte index 7
7

julia> validate_fasta(">hello\nTAGAGA\nTAGAG\n") # nothing; it matches

julia> eval(generate_io_validator(:validate_io, fasta_regex));

julia> validate_io(IOBuffer(">hello\nTAGAGA\n"))

julia> validate_io(IOBuffer(">helX"))
(0x58, (1, 5))

julia> validate_io(IOBuffer(">hello\n\n"))
(0x0a, (3, 0))

julia> validate_io(IOBuffer(">hello\nAC"))
(nothing, (2, 2))

julia> make_tokenizer(
    [
        re"\(",
        re"\)",
        re",",
        re"\"",
        re" +",
        re"[a-zA-Z]+"
    ]
) |> eval

machine = let
    header = onexit!(onenter!(re"[a-z]+", :mark_pos), :header)
    seqline = onexit!(onenter!(re"[ACGT]+", :mark_pos), :seqline)
    record = onexit!(re">" * header * '\n' * seqline * rep('\n' * seqline), :record)
    compile(opt(record) * rep('\n' * record) * rep(re"\n"))
end;
ERROR: Ambiguous NFA.

machine = let
    a = onenter!(re"XY", :a)
    b = onenter1(re"XZ", :b)
    compile('A' * (a | b))
end;

machine = let
    a = precond!(onenter!(re"XY", :a), :test)
    b = precond!(onenter!(re"XZ", :b), :test; bool=false) # <==
    compile('A' * (a | b))
end
# 上述代码中, precond!(regex, label)等效于precond!(regex, label; when=:enter, bool=true)
# 即仅当`:test`为`false`时, 再解析`re"XZ"`

struct Seq
    name::String
    seq::String
end

machine = let
    header = onexit!(onenter!(re"[a-z]+", :mark_pos), :header)
    seqline = onexit!(onenter!(re"[ACGT]+", :mark_pos), :seqline)
    record = onexit!(re">" * header * '\n' * rep1(seqline * '\n'), :record)
    compile(rep(record))
end
@assert machine isa Automa.Machine

actions = Dict(
    :mark_pos => :(pos = p),
    :header => :(header = String(data[pos:p-1])),
    :seqline => :(append!(buffer, data[pos:p-1])),
    :record => :(push!(seqs, Seq(header, String(buffer))))
);

@eval function parse_fasta(data)
    pos = 0
    buffer = UInt8[]
    seqs = Seq[] # 初始化seqs
    header = ""  # 初始化header
    $(generate_code(machine, actions)) # action dict生成相关代码
    return seqs
end

parse_fasta(">abc\nTAGA\nAAGA\n>header\nAAAG\nGGCG\n")

struct SeqPos
    offset::Int
    hlen::Int32
    slen::Int32
end

actions = Dict(
    :mark_pos => :(pos = p),
    :header => :(hlen = p - pos),
    :seqline => :(slen += p - pos),
    :record => quote
        seqpos = SeqPos(offset, hlen, slen)
        nseq += 1
        seqs[nseqs] = seqpos
        offset += hlen + slen
        slen = 0
    end
);
@assert actions isa Dict

@eval function parse_fasta(data)
    pos = slen = hlen = offset = nseqs = 0
    seqs = Vector{SeqPos}(undef, 400000)
    $(generate_code(machine, actions))
    return seqs
end

# generate_reader生成的代码示例:
function { function name }(stream::TranscodingStream, { args... })
    { init code }

    @label __exec__

    p = current buffer position
    p_end = final buffer position

    # the eof call below will first flush any used data from buffer,
    # then load in new data, before checking if it's really eof.
    is_eof = eof(stream)
    execute normal automa parsing of the buffer
    update buffer position to match p

    { loop code }

    if cs < 0 # meaning: erroneous input or erroneous EOF
        { error code }
    end

    if machine errored or reached EOF
        @label __return__
        { return code }
    end
    @goto __exec__
end

return_code = :(iszero(cs))
error_code = :(@goto __return__)
eval(generate_reader(:validate_fasta, machine; returncode=return_code, errorcode=error_code))

io = NoopStream(IOBuffer(">a\nTAG\nTA\n>bac\nG\n"))
validate_fasta(io) # true
validate_fasta(NoopStream(IOBuffer("random data"))) # false

struct Seq
    name::String
    seq::String
end

machine = let
    header = onexit!(onenter!(re"[a-z]+", :mark_pos), :header)
    seqline = onexit!(onenter!(re"[ACGT]+", :mark_pos), :seqline)
    record = onexit!(re">" * header * '\n' * rep1(seqline * '\n'), :record)
    compile(rep(record))
end
@assert machine isa Automa.Machine

# add @escape
actions = Dict{Symbol, Expr}(
    :mark_pos => :(pos = p),
    :header => :(header = String(data[pos:p-1])),
    :seqline => :(append!(seqbuffer, data[pos:p-1])),

    :record => quote
        seq = Seq(header, String(seqbuffer))
        found_sequence = true
        # 如果不是文件末尾, 就重新设置p = p-1
        p -= !(is_eof && p > p_end)
        @escape
    end
)

@assert actions isa Dict

generate_reader(
    :read_record,
    machine;
    actions=actions,
    initcode=quote
        seqbuffer = UInt8[]
        pos = 0
        found_sequence = false
        header = ""
    end,
    loopcode=quote
        if (is_eof && p > p_end) || found_sequence
            @goto __return__
        end
    end,
    returncode=:(found_sequence ?  seq : nothing)
) |> eval

io = NoopStream(IOBuffer(">\nT\n>tag\nGAGA\nTATA\n"));
read_record(io) # return the 1st seq record
read_record(io) # return the 2nd seq record
read_record(io) # return nothing

actions = Dict{Symbol, Expr}(
    :mark_position => :(@mark),
    :header => :(header = String(data[@markpos():p-1])),
    :seqline => :(append!(buffer, data[@markpos():p-1]))，

    # 其他代码略
)

struct Seq
    name::String
    seq::String
end

machine = let
    header = onexit!(onenter!(re"[a-z]+", :mark_pos), :header)
    seqline = onexit!(onenter!(re"[ACGT]+", :mark_pos), :seqline)
    record = onexit!(re">" * header * '\n' * rep1(seqline * '\n'), :record)
    compile(rep(record))
end
@assert machine isa Automa.Machine

# 配置Reader类型
mutable struct Reader{S <: TranscodingStream}
    io::S
    automa_stat::Int
end

Reader(io::TranscodingStream) = Reader{typeof(io)}(io, 1)
Reader(io::IO) = Reader(NoopStream(io))

actions = Dict{Symbol, Expr}(
    :mark_pos => :(@mark),
    :header => :(header = String(data[@markpos():p-1])),
    :seqline => :(append!(seqbuffer, data[@markpos():p-1])),
    :record => quote
        seq = Seq(header, String(seqbuffer))
        found_sequence = true
        @escape
    end
)

generate_reader(
    :read_record,
    machine;
    actions=actions,
    arguments=(:(start_state::Int),),
    initcode=quote
        seqbuffer = UInt8[]
        found_sequence = false
        header = ""
        cs = start_state
    end,
    loopcode=quote
        if (is_eof && p > p_end) || found_sequence
            @goto __return__
        end
    end,
    returncode=:(found_sequence ? (cs, seq) : throw(EOFError()))
) |> eval

function read_record(reader::Reader)
    (cs, seq) = read_record(reader.io, reader.automa_state)
    reader.automa_state = cs
    return seq
end

julia> reader = Reader(IOBuffer(">a\nT\n>tag\nGAG\nATATA\n"));

julia> read_record(reader)
Seq("a", "T")

julia> read_record(reader)
Seq("tag", "GAGATATA")

julia> read_record(reader)
ERROR: EOFError: read end of file

module BioStockholm

# 载入依赖包和依赖函数
import Automa
using Automa: @re_str, onenter!, onexit!
const re = Automa.RegExp
using OrderedCollections: OrderedDict

# 设置输出的数据结构/函数
export MSA

"""
    MSA{T}

Stockholm format for multiple sequence alignment with
annotations. Sequence data is of type `Vector{T}`.

## Fields

seq: seqname => seqdata
GF : per_file_feature => text
GS : seqname => per_seq_feature => text
GC : per_file_feature => seqdata
GR : seqname => per_seq_feature => seqdata

"""
Base.@kwdef struct MSA{T}
    # seqname => seqdata
    seq :: OrderedDict{String, Vector{T}} =
        OrderedDict{String, Vector{T}}()
    # per_file_feature => text
    GF  :: OrderedDict{String, String} =
        OrderedDict{String, String}()
    # seqname => per_seq_feature => text
    GS  :: OrderedDict{String, OrderedDict{String, String}} =
        OrderedDict{String, OrderedDict{String, String}}()
    # per_file_feature => seqdata
    GC  :: OrderedDict{String, Vector{T}} =
        OrderedDict{String, Vector{T}}()
    # seqname => per_seq_feature => seqdata
    GR  :: OrderedDict{String, OrderedDict{String, Vector{T}}} =
        OrderedDict{String, OrderedDict{String, Vector{T}}}()

    function MSA{T}(seq, GF, GS, GC, GR) where {T}
        if valtype(seq) === String
            seq = OrderedDict(name => T.(collect(s)) for (name,s) in seq)
        end
        if valtype(GC) === String
            GC = OrderedDict(feat => T.(collect(s)) for (feat,s) in GC)
        end
        if valtype(valtype(GR)) === String
            GR = OrderedDict(name => OrderedDict(feat => T.(collect(s)) for (feat,s) in f2s) for (name,f2s) in GR)
        end
        return new{T}(seq, GF, GS, GC, GR)
    end
end

# 下面是定义base.==和base.show, 与解析器无关, 这里就不粘贴了
function Base.:(==)(s1::MSA, s2::MSA)
    # pass
end

function Base.show(io::IO, mime::MIME"text/plain", msa::MSA)
    # pass
end

# 定义machine

const stockholm_machine = let
    nl      = re"\r?\n"         # 支持回车符的新行RE
    ws      = re"[ \t]+"        # 分隔符: 空格/制表符
    feature = re"[^ \t\r\n]+"   # 分隔符开头的行
    # TODO: slash '/' prohibited at beginning of seqname
    seqname = re"[^#/ \t\r\n][^ \t\r\n]*"   # 匹配序列名, 开头不能是#,/和分隔符, 后边允许是/
    text    = re"[^ \t\r\n][^\r\n]*"        
    seqdata = re"[^ \t\r\n]+"               

    line_header = re"# STOCKHOLM 1\.0" * nl         # 头文件
    line_end    = re"//" * nl                       # 记录结束
    line_GF     = re"#=GF" * ws * feature * ws * text * nl
    line_GC     = re"#=GC" * ws * feature * ws * seqdata * nl   
    line_GS     = re"#=GS" * ws * seqname * ws * feature * ws * text * nl
    line_GR     = re"#=GR" * ws * seqname * ws * feature * ws * seqdata * nl
    line_seq    = seqname * ws * seqdata * nl
    line_empty  = re"[ \t]*" * nl

    # 封装成整个条目的RE
    stockholm = (
        re.rep(line_empty) # 空行分配给后一个条目, 而不是前一个条目
        * line_header
        * re.rep(line_GF | line_GC | line_GS | line_GR | line_seq | line_empty)
        * line_end
    )

    onenter!(nl, :countline)
    onenter!(feature, :enter_feature)
    onexit!(feature, :feature)
    onenter!(seqname, :enter_seqname)
    onexit!(seqname, :seqname)
    onenter!(text, :enter_text)
    onexit!(text, :text)
    onenter!(seqdata, :enter_seqdata)
    onexit!(seqdata, :seqdata)
    onexit!(line_GF, :line_GF)
    onexit!(line_GC, :line_GC)
    onexit!(line_GS, :line_GS)
    onexit!(line_GR, :line_GR)
    onexit!(line_seq, :line_seq)

    Automa.compile(stockholm)
end

# 定义Actions
const stockholm_actions = Dict(
    :countline => :(linenum += 1), # 记录行号

    :enter_feature => :(mark = p),
    :enter_seqname => :(mark = p),
    :enter_text    => :(mark = p),
    :enter_seqdata => :(mark = p),

    :feature => :(feature = mark == 0 ? "" : String(data[mark:p-1]); mark = 0),
    :seqname => :(seqname = mark == 0 ? "" : String(data[mark:p-1]); mark = 0),
    :text    => :(text = mark == 0 ? "" : String(data[mark:p-1]); mark = 0),
    :seqdata => :(seqdata = mark == 0 ? T[] : T.(collect(data[mark:p-1])); mark = 0),

    :line_GF => quote
        if haskey(gf_records, feature)
            gf_records[feature] *= " " * text
        else
            gf_records[feature] = text
        end
    end,
    # 读取记录存dict
    :line_GC => :(
        # gc_records[feature] = get(gc_records, feature, "") * seqdata
        gc_records[feature] = append!(get(gc_records, feature, T[]), seqdata)
    ),
    :line_GS => quote
        if haskey(gs_records, seqname)
            if haskey(gs_records[seqname], feature)
                gs_records[seqname][feature] *= " " * text
            else
                gs_records[seqname][feature] = text
            end
        else
            gs_records[seqname] = OrderedDict(feature => text)
        end
    end,
    :line_GR => quote
        if haskey(gr_records, seqname)
            # gr_records[seqname][feature] = get(gr_records[seqname], feature, "") * seqdata
            gr_records[seqname][feature] = append!(get(gr_records[seqname], feature, T[]), seqdata)
        else
            gr_records[seqname] = OrderedDict(feature => seqdata)
        end
    end,
    :line_seq => :(
        sequences[seqname] = append!(get(sequences, seqname, T[]), seqdata)
    ),
)

# 派发Base.read和Base.write支持MSA的读写
Base.read(io::IO, ::Type{MSA{T}}) where {T} = parse(MSA, read(io, String))
Base.read(io::IO, ::Type{MSA}) = read(io, MSA{Char})
Base.read(filepath::AbstractString, ::Type{MSA{T}}) where {T} = open(filepath) do io
    read(io, MSA)
    end
Base.read(filepath::AbstractString, ::Type{MSA}) = read(filepath, MSA{String}) 

Base.write(io::IO, msa::MSA) = print(io, msa)
Base.write(filepath::AbstractString, msa::MSA) =
    open(filepath, "w") do io
        print(io, msa)
    end

# 派发Base.parse 支持 MSA
Base.parse(::Type{MSA}, data::Union{String, Vector{UInt8}}) = 
    parse_stockholm(Char, data)

Base.parse(::Type{MSA{T}}, data::Union{String, Vector{UInt8}}) where {T} =
    parse_stockholm(T, data)

const context = Automa.CodeGenContext(generator=:goto)
@eval function parse_stockholm(::Type{T}, data::Union{String, Vector{UInt8}}) where {T}
    # variables for the action code
    sequences  = OrderedDict{String,Vector{T}}()                        # seqname => seqdata
    gf_records = OrderedDict{String,String}()                           # feature => text
    gc_records = OrderedDict{String,Vector{T}}()                        # feature => seqdata
    gs_records = OrderedDict{String,OrderedDict{String,String}}()       # seqname => feature => text
    gr_records = OrderedDict{String,OrderedDict{String,Vector{T}}}()    # seqname => feature => seqdata
    linenum = 1
    mark = 0
    seqname = ""
    feature = ""
    text = ""
    seqdata = ""

    # init vars for state machine
    $(Automa.generate_init_code(context, stockholm_machine))
    p_end = p_eof = lastindex(data)

    # main loop over input data
    $(Automa.generate_exec_code(context, stockholm_machine, stockholm_actions))

    if cs != 0
        error("failed to parse on line ", linenum)
    end

    return MSA{T}(; seq=sequences, GF=gf_records, GC=gc_records,
                  GS=gs_records, GR=gr_records)
end

# 后边还有派发Base.print(), 与解析器无关, 就不粘贴了

end # module BioStockholm